{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "70ea1792-6f21-4a44-ae96-ffb4d65e894d",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "mat1 and mat2 shapes cannot be multiplied (64x64 and 1922x64)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 58\u001b[0m\n\u001b[1;32m     55\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m D\n\u001b[1;32m     57\u001b[0m \u001b[38;5;66;03m# 执行字典学习（计算字典D）\u001b[39;00m\n\u001b[0;32m---> 58\u001b[0m D \u001b[38;5;241m=\u001b[39m \u001b[43mdictionary_learning_torch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpatches_tensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mT\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_components\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m64\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     60\u001b[0m \u001b[38;5;66;03m# 4. 可视化字典集合（学习到的字典原型）\u001b[39;00m\n\u001b[1;32m     61\u001b[0m dictionary \u001b[38;5;241m=\u001b[39m D\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mnumpy()  \u001b[38;5;66;03m# 转移到CPU并转换为numpy数组\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[1], line 48\u001b[0m, in \u001b[0;36mdictionary_learning_torch\u001b[0;34m(X, n_components, n_iter)\u001b[0m\n\u001b[1;32m     44\u001b[0m D \u001b[38;5;241m=\u001b[39m D \u001b[38;5;241m/\u001b[39m D\u001b[38;5;241m.\u001b[39mnorm(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, keepdim\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)  \u001b[38;5;66;03m# 正则化字典\u001b[39;00m\n\u001b[1;32m     46\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(n_iter):\n\u001b[1;32m     47\u001b[0m     \u001b[38;5;66;03m# 计算系数矩阵 H\u001b[39;00m\n\u001b[0;32m---> 48\u001b[0m     H \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmatmul\u001b[49m\u001b[43m(\u001b[49m\u001b[43mD\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mX\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mT\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# D^T * X\u001b[39;00m\n\u001b[1;32m     49\u001b[0m     H \u001b[38;5;241m=\u001b[39m H \u001b[38;5;241m/\u001b[39m (H\u001b[38;5;241m.\u001b[39mnorm(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, keepdim\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1e-6\u001b[39m)  \u001b[38;5;66;03m# 正则化系数矩阵\u001b[39;00m\n\u001b[1;32m     51\u001b[0m     \u001b[38;5;66;03m# 更新字典 D\u001b[39;00m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: mat1 and mat2 shapes cannot be multiplied (64x64 and 1922x64)"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.decomposition import DictionaryLearning\n",
    "\n",
    "# 1. 加载并转换为灰度图\n",
    "target = 1\n",
    "data_name = ['0618', '0854', '1066'][target - 1]\n",
    "image = cv2.imread(f'../input_data/{data_name}.png')  # 加载RGB图像\n",
    "\n",
    "# 检查图像是否正确加载\n",
    "if image is None:\n",
    "    raise ValueError(f\"Error loading image: ../input_data/{data_name}.png\")\n",
    "\n",
    "# 转换为灰度图像\n",
    "gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# 调整图像大小以确保一致性（假设图像大小为500x250）\n",
    "gray_image = cv2.resize(gray_image, (250, 500))\n",
    "\n",
    "# 2. 使用字典学习提取特征\n",
    "patch_size = 8  # 假设每个块为8x8像素\n",
    "patches = []\n",
    "\n",
    "# 将图像分成多个8x8小块\n",
    "for y in range(0, gray_image.shape[0] - patch_size, patch_size):\n",
    "    for x in range(0, gray_image.shape[1] - patch_size, patch_size):\n",
    "        patch = gray_image[y:y + patch_size, x:x + patch_size]\n",
    "        patches.append(patch.flatten())  # 将块展平为一个特征向量\n",
    "\n",
    "# 转换为numpy数组\n",
    "patches = np.array(patches)\n",
    "\n",
    "# 将数据转换为PyTorch张量并移到GPU\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "patches_tensor = torch.tensor(patches).float().to(device)\n",
    "\n",
    "# 3. 使用GPU加速的字典学习（基于SVD）\n",
    "def dictionary_learning_torch(X, n_components=64, n_iter=100):\n",
    "    # 使用PyTorch实现字典学习（基于SVD）\n",
    "    m, n = X.shape\n",
    "    D = torch.randn(n_components, m, device=X.device)  # 字典D的大小应为 (n_components, m)，而不是 (n_components, n)\n",
    "    D = D / D.norm(dim=1, keepdim=True)  # 正则化字典\n",
    "\n",
    "    for _ in range(n_iter):\n",
    "        # 计算系数矩阵 H\n",
    "        H = torch.matmul(D, X.T)  # D^T * X\n",
    "        H = H / (H.norm(dim=0, keepdim=True) + 1e-6)  # 正则化系数矩阵\n",
    "\n",
    "        # 更新字典 D\n",
    "        X_hat = torch.matmul(D.T, H)  # D * H\n",
    "        D = X_hat / X_hat.norm(dim=1, keepdim=True)\n",
    "\n",
    "    return D\n",
    "\n",
    "# 执行字典学习（计算字典D）\n",
    "D = dictionary_learning_torch(patches_tensor.T, n_components=64)\n",
    "\n",
    "# 4. 可视化字典集合（学习到的字典原型）\n",
    "dictionary = D.cpu().detach().numpy()  # 转移到CPU并转换为numpy数组\n",
    "\n",
    "# 将字典中的每个原型reshape为图像块大小（8x8）\n",
    "fig, axes = plt.subplots(8, 8, figsize=(8, 8))\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    ax.imshow(dictionary[i].reshape(patch_size, patch_size), cmap='gray')\n",
    "    ax.axis('off')\n",
    "plt.suptitle(\"Dictionary Components (Learned Patches)\")\n",
    "plt.show()\n",
    "\n",
    "# 输出字典特征矩阵\n",
    "print(\"Dictionary Learning Results:\")\n",
    "print(\"Components (Shape):\", dictionary.shape)\n",
    "\n",
    "# 最终输出是1922x64的特征矩阵\n",
    "print(\"Output Feature Matrix Shape: \", patches.shape[0], \"x\", dictionary.shape[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b08de2-a165-4ff8-90ff-efac30a0da4b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
